import json
from tqdm import tqdm
from chat import chat_ark as chat_api
from process import get_json_from_llm

# 示例来自于 https://infinigence.feishu.cn/docx/SDdwdsNSto5yjcxOBk9cuTJhnag
create_template = """你是一名专业的智能体agent专家，请根据可调用的工具库 apis 和用户多轮对话，给出多轮对话中应使用的工具 api 名称和参数。

## 要求
1.用户的多轮对话中只有一个api被调用，给出工具名。
2.根据多轮对话内容，判断出在第几轮可以满足工具调用的参数需求，然后将api名称和参数加入该轮结果。
3.输出 json 格式：
```json
```

示例：
工具库 apis:
{example['apis']}

多轮对话:
{example['user_messages]}

输出：
{example['targets]}

工具库 apis:
{apis}

多轮对话:
{user_messages}

输出："""


prompt_template = """你是一名专业的智能体agent专家，请根据可调用的工具库 apis 和用户多轮对话，给出多轮对话中应使用的工具 api 名称和参数。

## 要求
1.用户的多轮对话中只有一个api被调用，给出工具名。
2.根据多轮对话内容，判断出在第几轮可以满足工具调用的参数需求，然后将api名称和参数加入该轮结果。

工具库 apis:
{apis}

多轮对话:
{user_messages}

输出："""


res = []
with open('data/合成数据初赛赛题.jsonl','r', encoding="utf-8") as f:
    for line in tqdm(f):
        line = line.strip()
        if line == '':
            continue
        d = json.loads(line)

        apis = d['apis']
        user_messages = d['user_messages']
        
        create_prompt = create_template.replace('{apis}', json.dumps(apis, ensure_ascii=False))
        create_prompt = create_prompt.replace('{user_messages}', json.dumps(user_messages, ensure_ascii=False))
        res = chat_api(create_prompt)
        targets = get_json_from_llm(res)

        prompt = prompt_template.replace('{apis}', json.dumps(apis, ensure_ascii=False))
        prompt = prompt.replace('{user_messages}', json.dumps(user_messages, ensure_ascii=False))

        res.append({
            'instruction': prompt,
            'input': '',
            'output': json.dumps(targets, ensure_ascii=False)
        })
with open('data/train_atom.json', 'w', encoding="utf-8") as f:
    json.dump(res, f, ensure_ascii=False, indent=4)